#!/usr/bin/env python3
"""
依赖安全扫描脚本

使用 pip-audit 和 safety 工具扫描项目依赖中的安全漏洞
"""
import json
import subprocess
import sys
from pathlib import Path

def run_command(cmd: list) -> tuple[int, str]:
    """运行命令并返回结果"""
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=False
        )
        return result.returncode, result.stdout + result.stderr
    except Exception as e:
        return 1, str(e)

def scan_with_pip_audit(requirements_file: str) -> bool:
    """使用 pip-audit 扫描依赖"""
    print(f"\n🔍 使用 pip-audit 扫描 {requirements_file}...")
    
    # 先运行扫描获取结果
    cmd = [
        "python", "-m", "pip_audit",
        "--requirement", requirements_file,
        "--format", "json"
    ]
    
    returncode, output = run_command(cmd)
    
    # 保存和处理报告
    report_file = f"security-reports/pip-audit-{Path(requirements_file).stem}.json"
    try:
        if output:
            data = json.loads(output)
            with open(report_file, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            
            # 显示漏洞详情
            if data.get("vulnerabilities"):
                print(f"⚠️  发现 {len(data['vulnerabilities'])} 个漏洞:")
                for vuln in data["vulnerabilities"]:
                    package = vuln.get("package", "未知")
                    version = vuln.get("installed_version", "未知")
                    vuln_id = vuln.get("id", "未知")
                    summary = vuln.get("description", "无描述")[:80]
                    fix_versions = vuln.get("fix_versions", [])
                    
                    print(f"  📦 {package} v{version} - {vuln_id}")
                    print(f"     {summary}...")
                    if fix_versions:
                        print(f"     🔧 修复版本: {', '.join(fix_versions)}")
                return False
            else:
                print("✅ 未发现漏洞")
                return True
        else:
            print("✅ pip-audit 扫描完成，未发现问题")
            return True
    except json.JSONDecodeError:
        print(f"❌ 解析 pip-audit 输出失败")
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump({"error": "JSON解析失败", "raw_output": output}, f, indent=2, ensure_ascii=False)
        return False
    except Exception as e:
        print(f"❌ 处理 pip-audit 结果失败: {e}")
        return False

def scan_with_safety(requirements_file: str) -> bool:
    """使用 safety 扫描依赖"""
    print(f"\n🔍 使用 safety 扫描 {requirements_file}...")
    
    cmd = [
        "python", "-m", "safety",
        "check",
        "--requirement", requirements_file,
        "--json"
    ]
    
    returncode, output = run_command(cmd)
    
    # 保存和处理报告
    report_file = f"security-reports/safety-{Path(requirements_file).stem}.json"
    try:
        if output:
            # Safety可能返回非JSON格式的错误信息
            if output.strip().startswith('[') or output.strip().startswith('{'):
                data = json.loads(output)
                with open(report_file, 'w', encoding='utf-8') as f:
                    json.dump(data, f, indent=2, ensure_ascii=False)
                
                # 显示漏洞详情
                if isinstance(data, list) and data:
                    print(f"⚠️  发现 {len(data)} 个安全问题:")
                    for issue in data:
                        package = issue.get("package", "未知")
                        version = issue.get("installed_version", "未知")
                        vuln_id = issue.get("vulnerability_id", "未知")
                        advisory = issue.get("advisory", "无描述")[:80]
                        
                        print(f"  📦 {package} v{version} - {vuln_id}")
                        print(f"     {advisory}...")
                    return False
                else:
                    print("✅ 未发现安全问题")
                    return True
            else:
                # 非JSON输出，可能是错误信息
                with open(report_file, 'w', encoding='utf-8') as f:
                    json.dump({"error": "非JSON输出", "raw_output": output}, f, indent=2, ensure_ascii=False)
                if "No known security vulnerabilities" in output:
                    print("✅ 未发现安全问题")
                    return True
                else:
                    print(f"⚠️ Safety 输出: {output}")
                    return False
        else:
            print("✅ safety 扫描完成，未发现问题")
            with open(report_file, 'w', encoding='utf-8') as f:
                json.dump({"vulnerabilities": [], "message": "无漏洞"}, f, indent=2, ensure_ascii=False)
            return True
    except json.JSONDecodeError:
        print(f"❌ 解析 safety 输出失败")
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump({"error": "JSON解析失败", "raw_output": output}, f, indent=2, ensure_ascii=False)
        return False
    except Exception as e:
        print(f"❌ 处理 safety 结果失败: {e}")
        return False

def main():
    """主函数"""
    print("🔒 开始依赖安全扫描...")
    
    # 确保报告目录存在
    Path("security-reports").mkdir(exist_ok=True)
    
    # 要扫描的依赖文件
    requirements_files = [
        "requirements/base.txt",
        "requirements/development.txt",
        "requirements/testing.txt",
        "requirements/production.txt"
    ]
    
    all_passed = True
    
    for req_file in requirements_files:
        if not Path(req_file).exists():
            print(f"⚠️ 跳过不存在的文件: {req_file}")
            continue
            
        # 使用两个工具扫描
        pip_audit_passed = scan_with_pip_audit(req_file)
        safety_passed = scan_with_safety(req_file)
        
        if not (pip_audit_passed and safety_passed):
            all_passed = False
    
    if all_passed:
        print("\n✅ 所有依赖安全扫描通过！")
        sys.exit(0)
    else:
        print("\n❌ 发现安全漏洞，请检查报告并更新相关依赖")
        sys.exit(1)

if __name__ == "__main__":
    main()
